#!/usr/bin/python

import numpy as np
import scipy
import frame_convert

PLYUNIT = 1000

class pointcloud(object):    
    def __init__(self, depth = None, color = None):
        '''
        Init function takes aligned, filtered, smoothed depth and color matrices
        and convert them to a point cloud.
        Assume color is a 3-D array with last element represents RGB
        (R = 0, G = 1, B = 2).
        This function will ignore all points that has a negtive depth.
        '''
        
        if depth == None:
            self.vertex = np.zeros((0, 3))
            self.color = np.zeros((0, 3))
            return
        
        assert depth.ndim == 2
        if __debug__:
	    if color != None:
	        if not color.ndim == 3:
	            raise AssertionError
	        (x, y, z) = color.shape
	        if not z == 3:
	            raise AssertionError
	        if not depth.shape == (x, y):
	            raise AssertionError

        # get dimensions
        (rows, cols) = depth.shape
        
        # if color not given, assume grey (0x808080)
        if color == None:
	    color = np.empty((rows, cols, 3))
	    color[:,:,:] = 0x80
        
        # prepare input for world convert
        vertex = np.empty((rows * cols, 3))
        # x = 0 ... col-1 0 ... col-1 ... repeat row times
        vertex[:,0] = np.tile(np.arange(cols), rows).astype(float)
        # y = 0 0 0 ... repeat col times 1 1 1 ... repeat row times
        vertex[:,1] = np.tile(np.arange(rows), (cols, 1)).T.flatten()
        vertex[:,2] = depth.flatten()
        
        # prepare color matrix
        colortemp = color.reshape((rows * cols, 3))
        
        # eliminate negtive depth points
        index = vertex[:,2] >= 0
        vertex = vertex[index, :]
        colortemp = colortemp[index, :]
        
        # do world convert
        vertex = frame_convert.depth_world(vertex.T).T
        
        self.vertex = vertex
        self.color = colortemp
        return
    
    def __add__(self, other):
        new = pointcloud()
        new.vertex = np.concatenate((self.vertex, other.vertex), 0)
        new.color = np.concatenate((self.color, other.color), 0)
        return new
    
    def __iadd__(self, other):
        self.vertex = np.concatenate((self.vertex, other.vertex), 0)
        self.color = np.concatenate((self.color, other.color), 0)
        return self
        
    def clipCube(self, left=None, right=None, top=None,\
        bottom=None, front=None, back=None):
        '''
        Clip the point cloud into a cube.
        All parameter in real unit(meters), there are from the camera's point
        of view, e.g. left = x_neative, top = y_negative, front = z_negative.
        If any parameter is None, ignore that boundary.
        '''
        
        if left != None:
            index = self.vertex[:,0] >= left
            self.vertex = self.vertex[index, :]
            self.color = self.color[index, :]
        
        if right != None:
            index = self.vertex[:,0] <= right
            self.vertex = self.vertex[index, :]
            self.color = self.color[index, :]
        
        if top != None:
            index = self.vertex[:,1] >= top
            self.vertex = self.vertex[index, :]
            self.color = self.color[index, :]
        
        if bottom != None:
            index = self.vertex[:,1] <= bottom
            self.vertex = self.vertex[index, :]
            self.color = self.color[index, :]
        
        if front != None:
            index = self.vertex[:,2] >= front
            self.vertex = self.vertex[index, :]
            self.color = self.color[index, :]
        
        if back != None:
            index = self.vertex[:,2] <= back
            self.vertex = self.vertex[index, :]
            self.color = self.color[index, :]
        return
    
    def toPly(self, fout):
        '''
        This function export a pointcloud instance to a ply format file.
        '''
        vertexCount = self.vertex.shape[0]
        # prepare header
        fout.write("ply\n")
        fout.write("format ascii 1.0\n")
        fout.write("element vertex %d\n" % vertexCount)
        fout.write("property float x\n")
        fout.write("property float y\n")
        fout.write("property float z\n")
        fout.write("property uchar red\n")
        fout.write("property uchar green\n")
        fout.write("property uchar blue\n")
        fout.write("end_header\n")
        # write vertices
        for i in range(vertexCount):
            fout.write("%f %f %f %d %d %d\n" % \
                        (self.vertex[i,0]*PLYUNIT, \
                        self.vertex[i,1]*PLYUNIT, \
                        self.vertex[i,2]*PLYUNIT, \
                        self.color[i,0], self.color[i,1], self.color[i,2]))
        return

if __name__ == "__main__":
    import scipy.io
    import sys
    
    argc = len(sys.argv)
    if argc < 3:
        print "Usage: pointcloud outputfile inputfile"
        sys.exit(1)
    
    mat = scipy.io.loadmat(sys.argv[2])
    depth = mat['depth']
    color = mat['rgb']
    
    fout = open(sys.argv[1], 'w')
    pc = pointcloud(depth, color)
    pc.toPly(fout)
    fout.close()

